import os
import pandas as pd
import numpy as np
import chromadb
from chromadb.utils import embedding_functions
from loguru import logger
import jieba
import re
from typing import List, Dict, Tuple, Optional

class MedicalKnowledgeBase:
    """医疗知识库类，用于存储和检索医疗知识"""
    
    def __init__(self, data_dir: str = None, collection_name: str = "medical_knowledge"):
        """初始化医疗知识库
        
        Args:
            data_dir: 医疗数据文件所在目录，默认为项目根目录下的data目录
            collection_name: 知识库集合名称
        """
        # 设置数据目录
        if data_dir is None:
            # 获取项目根目录
            root_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
            self.data_dir = os.path.join(root_dir, "data")
            logger.info(f"未指定数据目录，使用默认目录: {self.data_dir}")
        else:
            self.data_dir = data_dir
        
        logger.info(f"数据目录: {self.data_dir}")
        # 初始化向量数据库
        self.client = chromadb.Client()
        
        # 使用默认的embedding函数
        self.embedding_function = embedding_functions.DefaultEmbeddingFunction()
        
        # 创建或获取集合
        try:
            self.collection = self.client.get_or_create_collection(
                name=collection_name,
                embedding_function=self.embedding_function
            )
            logger.info(f"成功初始化知识库集合: {collection_name}")
        except Exception as e:
            logger.error(f"初始化知识库集合失败: {str(e)}")
            raise ValueError(f"知识库初始化失败: {str(e)}")
        
        # 记录已加载的数据文件
        self.loaded_files = set()
        
    def _preprocess_text(self, text: str) -> str:
        """预处理文本
        
        Args:
            text: 输入文本
            
        Returns:
            预处理后的文本
        """
        if not text or not isinstance(text, str):
            return ""
            
        # 去除多余空白字符
        text = re.sub(r'\s+', ' ', text.strip())
        # 分词并重新组合
        words = jieba.cut(text)
        return " ".join(words)
    
    def load_data(self, department: str = None) -> int:
        """加载指定科室的医疗数据
        
        Args:
            department: 科室名称，如果为None则加载所有科室数据
            
        Returns:
            加载的数据条目数量
        """
        # 获取数据文件列表
        data_files = [f for f in os.listdir(self.data_dir) if f.endswith('.csv')]
        
        if not data_files:
            logger.warning(f"在{self.data_dir}目录下未找到CSV数据文件")
            return 0
        
        # 如果指定了科室，只加载该科室的数据
        if department:
            data_files = [f for f in data_files if department in f]
            if not data_files:
                logger.warning(f"未找到{department}相关的数据文件")
                return 0
        logger.info(f"共找到{len(data_files)}个数据文件")
        total_loaded = 0
        
        for file_name in data_files:
            # 如果文件已经加载过，跳过
            if file_name in self.loaded_files:
                continue
                
            file_path = os.path.join(self.data_dir, file_name)
            try:
                # 读取CSV文件
                df = pd.read_csv(file_path)
                logger.warning( str(len(df))+"数据")
                
                # 确保数据包含问题和答案列
                if 'question' not in df.columns or 'answer' not in df.columns:
                    logger.warning(f"文件{file_name}缺少必要的列(question/answer)")
                    continue
                
                
                # 预处理数据
                documents = []
                metadatas = []
                ids = []
                
                for idx, row in df.iterrows():
                    question = self._preprocess_text(row['question'])
                    answer = row['answer']
                    
                    if not question or not answer:
                        continue
                    
                    doc_id = f"{file_name.split('.')[0]}_{idx}"
                    
                    documents.append(question)
                    metadatas.append({
                        "answer": answer,
                        "source": file_name,
                        "department": department or file_name.split('-')[0]
                    })
                    ids.append(doc_id)
                
                # 批量添加到向量数据库
                if documents:
                    self.collection.add(
                        documents=documents,
                        metadatas=metadatas,
                        ids=ids
                    )
                    
                    total_loaded += len(documents)
                    self.loaded_files.add(file_name)
                    logger.info(f"成功加载{file_name}，共{len(documents)}条数据")
                    
            except Exception as e:
                logger.error(f"加载{file_name}失败: {str(e)}")
        
        return total_loaded
    
    def search(self, query: str, department: str = None, top_k: int = 3) -> List[Dict]:
        """搜索相关医疗知识
        
        Args:
            query: 查询文本
            department: 限定科室范围
            top_k: 返回结果数量
            
        Returns:
            检索结果列表，每个元素包含答案和来源信息
        """
        # 预处理查询文本
        processed_query = self._preprocess_text(query)
        
        if not processed_query:
            logger.warning("查询文本为空")
            return []
        
        try:
            # 执行查询
            results = self.collection.query(
                query_texts=[processed_query],
                n_results=top_k,
                where={"department": department} if department else None
            )
            
            # 处理结果
            if not results["metadatas"] or not results["metadatas"][0]:
                return []
                
            retrieved_docs = []
            for metadata in results["metadatas"][0]:
                retrieved_docs.append({
                    "answer": metadata["answer"],
                    "source": metadata["source"],
                    "department": metadata["department"]
                })
                
            return retrieved_docs
            
        except Exception as e:
            logger.error(f"知识库检索失败: {str(e)}")
            return []
    
    def get_sources(self, results: List[Dict]) -> List[str]:
        """从检索结果中提取来源信息
        
        Args:
            results: 检索结果列表
            
        Returns:
            来源列表
        """
        if not results:
            return ["《医学知识库》"]
            
        sources = set()
        for item in results:
            if "source" in item:
                source_name = item["source"].split('.')[0]
                # 格式化来源名称
                if "-" in source_name:
                    dept = source_name.split('-')[0]
                    sources.add(f"《{dept}医学知识库》")
                else:
                    sources.add(f"《{source_name}》")
        
        return list(sources) if sources else ["《医学知识库》"]